package org.idea.qiyu.framework.apm.agent.core.dubbo;

import javassist.*;
import org.apache.dubbo.registry.support.AbstractRegistryFactory;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.RpcInvocation;
import org.idea.qiyu.common.constants.RequestConstants;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.UUID;

import static org.idea.qiyu.common.constants.RequestConstants.QIYU_BASE_REQUEST;

/**
 * @Author linhao
 * @Date created in 4:17 下午 2021/7/19
 */
public class DubboAgent {

    private static final String MOCK_CLUSTER_INVOKER = "org/apache/dubbo/rpc/cluster/support/wrapper/MockClusterInvoker";

    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (loader == null) {
                    return null;
                }
                if (!MOCK_CLUSTER_INVOKER.equals(className)) {
                    return null;
                }
                return buildDubboAgentClass(className.replace("/", "."));
            }
        }, true);
    }


    /**
     * 构建Dubbo的agent对象
     *
     * @param className
     * @return
     */
    public static byte[] buildDubboAgentClass(String className) {
        try {
            ClassPool classPool = ClassPool.getDefault();
            classPool.appendSystemPath();
            CtClass ctClass = classPool.get(className);
            CtMethod ctMethod = ctClass.getDeclaredMethod("invoke",
                    classPool.get(new String[]{
                            "org.apache.dubbo.rpc.Invocation"
                    }));
            CtMethod copyMethod = CtNewMethod.copy(ctMethod, ctClass, new ClassMap());
            String newMethodName = copyMethod.getName() + "$Agent";
            ctMethod.setName(newMethodName);
            copyMethod.setBody("{" +
                    "try{" +
                    "org.idea.qiyu.framework.apm.agent.core.dubbo.DubboAgent.begin($$);" +
                    " return " + newMethodName + "($$);" +
                    "}finally {" +
                    "org.idea.qiyu.framework.apm.agent.core.dubbo.DubboAgent.end();" +
                    "}" +
                    "}");
            ctClass.addMethod(copyMethod);
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static void doReport(Invocation invocation) {
        try {
            String qiyuBaseRequest = invocation.getAttachments().get(String.valueOf(RequestConstants.QIYU_BASE_REQUEST));

            org.idea.qiyu.common.dto.AgentCollectorDataDTO agentCollectorDataDTO = new org.idea.qiyu.common.dto.AgentCollectorDataDTO();
            agentCollectorDataDTO.setReportType((short) 1);
            agentCollectorDataDTO.setHeadTraceId(invocation.getAttachment(String.valueOf(QIYU_BASE_REQUEST)));
            agentCollectorDataDTO.setJson("");
            agentCollectorDataDTO.setTimestamp(System.currentTimeMillis());
            org.idea.qiyu.common.utils.ApmAgentUtil.doReport(agentCollectorDataDTO);
            System.out.println("日志上报");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void begin(Invocation invocation) {
        DubboStatisticsObject dubboStatisticsObject = DubboTraceSession.get();
        if (dubboStatisticsObject == null) {
            RpcInvocation rpcInvocation = (RpcInvocation) invocation;
            dubboStatisticsObject = new DubboStatisticsObject();
            dubboStatisticsObject.setArguments(rpcInvocation.getArguments());
            dubboStatisticsObject.setMethodName(rpcInvocation.getMethodName());
            dubboStatisticsObject.setAttachments(rpcInvocation.getAttachments());
            dubboStatisticsObject.setParameterTypes(rpcInvocation.getParameterTypes());
            dubboStatisticsObject.setBeginTime(System.currentTimeMillis());
            dubboStatisticsObject.setTraceId(UUID.randomUUID().toString());
            DubboTraceSession.set(dubboStatisticsObject);
        }
    }

    public static void end() {
        DubboStatisticsObject dubboStatisticsObject = DubboTraceSession.get();
        if (dubboStatisticsObject != null) {
            dubboStatisticsObject.setEndTime(System.currentTimeMillis());
            dubboStatisticsObject.setCount(dubboStatisticsObject.getEndTime() - dubboStatisticsObject.getBeginTime());
            System.out.println(dubboStatisticsObject);
        }
        DubboTraceSession.remove();
    }
}
